# This file is part of Androguard.
#
# Copyright (C) 2012, Anthony Desnos <desnos at t0t0.fr>
# All rights reserved.
#
# Androguard is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Androguard is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with Androguard.  If not, see <http://www.gnu.org/licenses/>.

from struct import unpack, pack

from androguard.core.androconf import warning, error, CONF, enable_colors, remove_colors, save_colors


def disable_print_colors():
    colors = save_colors()
    remove_colors()
    return colors


def enable_print_colors(colors):
    enable_colors(colors)

# Handle exit message
def Exit(msg):
    warning("Error : " + msg)
    raise ("oops")


def Warning(msg):
    warning(msg)


def _PrintBanner():
    print_fct = CONF["PRINT_FCT"]
    print_fct("*" * 75 + "\n")


def _PrintSubBanner(title=None):
    print_fct = CONF["PRINT_FCT"]
    if title == None:
        print_fct("#" * 20 + "\n")
    else:
        print_fct("#" * 10 + " " + title + "\n")


def _PrintNote(note, tab=0):
    print_fct = CONF["PRINT_FCT"]
    note_color = CONF["COLORS"]["NOTE"]
    normal_color = CONF["COLORS"]["NORMAL"]
    print_fct("\t" * tab + "%s# %s%s" % (note_color, note, normal_color) + "\n")

# Print arg into a correct format
def _Print(name, arg):
    buff = name + " "

    if type(arg).__name__ == 'int':
        buff += "0x%x" % arg
    elif type(arg).__name__ == 'long':
        buff += "0x%x" % arg
    elif type(arg).__name__ == 'str':
        buff += "%s" % arg
    elif isinstance(arg, SV):
        buff += "0x%x" % arg.get_value()
    elif isinstance(arg, SVs):
        buff += arg.get_value().__str__()

    print buff


def PrettyShow(basic_blocks):
    PrettyShow1(basic_blocks)


def PrettyShowEx(exceptions):
    if len(exceptions) > 0:
        CONF["PRINT_FCT"]("Exceptions:\n")
        for i in exceptions:
            CONF["PRINT_FCT"]("\t%s%s%s\n" % (CONF["COLORS"]["EXCEPTION"], i.show_buff(), CONF["COLORS"]["NORMAL"]))


def _PrintXRef(tag, items):
    print_fct = CONF["PRINT_FCT"]
    for i in items:
        print_fct("%s: %s %s %s %s\n" % (
            tag, i[0].get_class_name(), i[0].get_name(), i[0].get_descriptor(),
            ' '.join("%x" % j.get_idx() for j in i[1])))


def _PrintDefault(msg):
    print_fct = CONF["PRINT_FCT"]
    print_fct(msg)


def PrettyShow1(basic_blocks):
    idx = 0
    nb = 0

    offset_color = CONF["COLORS"]["OFFSET"]
    offset_addr_color = CONF["COLORS"]["OFFSET_ADDR"]
    instruction_name_color = CONF["COLORS"]["INSTRUCTION_NAME"]
    branch_false_color = CONF["COLORS"]["BRANCH_FALSE"]
    branch_true_color = CONF["COLORS"]["BRANCH_TRUE"]
    branch_color = CONF["COLORS"]["BRANCH"]
    exception_color = CONF["COLORS"]["EXCEPTION"]
    bb_color = CONF["COLORS"]["BB"]
    normal_color = CONF["COLORS"]["NORMAL"]
    print_fct = CONF["PRINT_FCT"]

    for i in basic_blocks:
        print_fct("%s%s%s : \n" % (bb_color, i.name, normal_color))
        instructions = i.get_instructions()
        for ins in instructions:
        #for ins in i.ins :

            notes = ins.get_notes()
            if notes != []:
                for note in notes:
                    _PrintNote(note, 1)

            print_fct("\t%s%-3d%s(%s%08x%s) " % (offset_color, nb, normal_color, offset_addr_color, idx, normal_color))
            print_fct("%s%-20s%s %s" % (instruction_name_color, ins.get_name(), normal_color, ins.get_output(idx)))

            op_value = ins.get_op_value()
            if ins == instructions[-1] and i.childs != []:
                print_fct(" ")
                # packed/sparse-switch
                if (op_value == 0x2b or op_value == 0x2c) and len(i.childs) > 1:
                    values = i.get_special_ins(idx).get_values()

                    print_fct("%s[ D:%s%s " % (branch_false_color, i.childs[0][2].name, branch_color))
                    print_fct(' '.join("%d:%s" % (values[j], i.childs[j + 1][2].name) for j in
                                       range(0, len(i.childs) - 1)) + " ]%s" % normal_color)
                else:
                    if len(i.childs) == 2:
                        print_fct("%s[ %s%s " % (branch_false_color, i.childs[0][2].name, branch_true_color))
                        print_fct(' '.join("%s" % c[2].name for c in i.childs[1:]) + " ]%s" % normal_color)
                    else:
                        print_fct("%s[ " % branch_color + ' '.join(
                            "%s" % c[2].name for c in i.childs) + " ]%s" % normal_color)

            idx += ins.get_length()
            nb += 1

            print_fct("\n")

        if i.exception_analysis != None:
            print_fct("\t%s%s%s\n" % (exception_color, i.exception_analysis.show_buff(), normal_color))

        print_fct("\n")


def method2dot(mx):
    """
        Export analysis method to dot format 

        @param mx : MethodAnalysis object

        @rtype : dot format buffer
    """

    vm = mx.get_vm()
    buff = ""
    for i in mx.basic_blocks.get():
        val = "green"
        if len(i.childs) > 1:
            val = "red"
        elif len(i.childs) == 1:
            val = "blue"

        for j in i.childs:
            buff += "\"%s\" -> \"%s\" [color=\"%s\"];\n" % ( i.get_name(), j[-1].get_name(), val )
            if val == "red":
                val = "green"

        idx = i.start
        label = ""
        for ins in i.get_instructions():
            label += "%x %s\l" % (idx, vm.dotbuff(ins, idx))
            idx += ins.get_length()

        buff += "\"%s\" [color=\"lightgray\", label=\"%s\"]\n" % (i.get_name(), label)
    return buff


def method2format(output, _format="png", mx=None, raw=False):
    """
        Export method to a specific file format

        @param output : output filename
        @param _format : format type (png, jpg ...) (default : png)
        @param mx : specify the MethodAnalysis object
        @param raw : use directly a dot raw buffer
    """
    try:
        import pydot
    except ImportError:
        error("module pydot not found")

    buff = "digraph code {\n"
    buff += "graph [bgcolor=white];\n"
    buff += "node [color=lightgray, style=filled shape=box fontname=\"Courier\" fontsize=\"8\"];\n"

    if raw == False:
        buff += method2dot(mx)
    else:
        buff += raw

    buff += "}"

    d = pydot.graph_from_dot_data(buff)
    if d:
        getattr(d, "write_" + _format)(output)


def method2png(output, mx=None, raw=False):
    """
        Export method to a png file format

        @param output : output filename
        @param mx : specify the MethodAnalysis object
        @param raw : use directly a dot raw buffer
    """
    buff = raw
    if raw == False:
        buff = method2dot(mx)

    method2format(output, "png", mx, buff)


def method2jpg(output, mx=None, raw=False):
    """
        Export method to a jpg file format

        @param output : output filename
        @param mx : specify the MethodAnalysis object
        @param raw : use directly a dot raw buffer
    """
    buff = raw
    if raw == False:
        buff = method2dot(mx)

    method2format(output, "jpg", mx, buff)


class SV:
    """SV is used to handle more easily a value"""

    def __init__(self, size, buff):
        self.__size = size
        self.__value = unpack(self.__size, buff)[0]

    def _get(self):
        return pack(self.__size, self.__value)

    def __str__(self):
        return "0x%x" % self.__value

    def __int__(self):
        return self.__value

    def get_value_buff(self):
        return self._get()

    def get_value(self):
        return self.__value

    def set_value(self, attr):
        self.__value = attr


class SVs:
    """SVs is used to handle more easily a structure of different values"""

    def __init__(self, size, ntuple, buff):
        self.__size = size

        self.__value = ntuple._make(unpack(self.__size, buff))

    def _get(self):
        l = []
        for i in self.__value._fields:
            l.append(getattr(self.__value, i))
        return pack(self.__size, *l)

    def _export(self):
        return [x for x in self.__value._fields]

    def get_value_buff(self):
        return self._get()

    def get_value(self):
        return self.__value

    def set_value(self, attr):
        self.__value = self.__value._replace(**attr)

    def __str__(self):
        return self.__value.__str__()


def object_to_str(obj):
    if isinstance(obj, str):
        return obj
    elif isinstance(obj, bool):
        return ""
    elif isinstance(obj, int):
        return pack("<L", obj)
    elif obj == None:
        return ""
    else:
        #print type(obj), obj
        return obj.get_raw()


class MethodBC(object):
    def show(self, value):
        getattr(self, "show_" + value)()


class BuffHandle:
    def __init__(self, buff):
        self.__buff = buff
        self.__idx = 0

    def read_b(self, size):
        return self.__buff[self.__idx: self.__idx + size]

    def read(self, size):
        if isinstance(size, SV):
            size = size.value

        buff = self.__buff[self.__idx: self.__idx + size]
        self.__idx += size

        return buff

    def end(self):
        return self.__idx == len(self.__buff)


class Buff:
    def __init__(self, offset, buff):
        self.offset = offset
        self.buff = buff

        self.size = len(buff)


class _Bytecode(object):
    def __init__(self, buff):
        try:
            import psyco

            psyco.full()
        except ImportError:
            pass

        self.__buff = buff
        self.__idx = 0

    def read(self, size):
        if isinstance(size, SV):
            size = size.value

        buff = self.__buff[self.__idx: self.__idx + size]
        self.__idx += size

        return buff

    def readat(self, off):
        if isinstance(off, SV):
            off = off.value

        return self.__buff[off:]

    def read_b(self, size):
        return self.__buff[self.__idx: self.__idx + size]

    def set_idx(self, idx):
        if isinstance(idx, SV):
            self.__idx = idx.value
        else:
            self.__idx = idx

    def get_idx(self):
        return self.__idx

    def add_idx(self, idx):
        self.__idx += idx

    def register(self, type_register, fct):
        self.__registers[type_register].append(fct)

    def get_buff(self):
        return self.__buff

    def length_buff(self):
        return len(self.__buff)

    def save(self, filename):
        fd = open(filename, "w")
        buff = self._save()
        fd.write(buff)
        fd.close()


def FormatClassToJava(input):
    """
       Transofmr a typical xml format class into java format

       @param input : the input class name
    """
    return "L" + input.replace(".", "/") + ";"


def FormatClassToPython(input):
    i = input[:-1]
    i = i.replace("/", "_")
    i = i.replace("$", "_")

    return i


def FormatNameToPython(input):
    i = input.replace("<", "")
    i = i.replace(">", "")
    i = i.replace("$", "_")

    return i


def FormatDescriptorToPython(input):
    i = input.replace("/", "_")
    i = i.replace(";", "")
    i = i.replace("[", "")
    i = i.replace("(", "")
    i = i.replace(")", "")
    i = i.replace(" ", "")
    i = i.replace("$", "")

    return i

####################### class/method/field export ########################
def ExportVMToPython(vm):
    """
        Export classes/methods/fields' names in the python namespace

        @param vm : a VM object (DalvikVMFormat, JVMFormat)
    """
    for _class in vm.get_classes():
        ### Class
        name = "CLASS_" + FormatClassToPython(_class.get_name())
        setattr(vm, name, _class)

        ### Methods
        m = {}
        for method in _class.get_methods():
            if method.get_name() not in m:
                m[method.get_name()] = []
            m[method.get_name()].append(method)

        for i in m:
            if len(m[i]) == 1:
                j = m[i][0]
                name = "METHOD_" + FormatNameToPython(j.get_name())
                setattr(_class, name, j)
            else:
                for j in m[i]:
                    name = "METHOD_" + FormatNameToPython(j.get_name()) + "_" + FormatDescriptorToPython(
                        j.get_descriptor())
                    setattr(_class, name, j)

        ### Fields
        f = {}
        for field in _class.get_fields():
            if field.get_name() not in f:
                f[field.get_name()] = []
            f[field.get_name()].append(field)

        for i in f:
            if len(f[i]) == 1:
                j = f[i][0]
                name = "FIELD_" + FormatNameToPython(j.get_name())
                setattr(_class, name, j)
            else:
                for j in f[i]:
                    name = "FIELD_" + FormatNameToPython(j.get_name()) + "_" + FormatDescriptorToPython(
                        j.get_descriptor())
                    setattr(_class, name, j)

